package cn.t.dynamic.switcher.configuration;

import cn.t.dynamic.switcher.router.DynamicDataSource;
import cn.t.dynamic.switcher.router.DynamicDataSourceContextHolder;
import cn.t.dynamic.switcher.type.DataSourceType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.annotation.Primary;
import org.springframework.util.CollectionUtils;

import javax.sql.DataSource;
import java.util.Map;

public class DynamicContextAware implements ApplicationContextAware {

    private static final Logger logger = LoggerFactory.getLogger(DynamicContextAware.class);

    /**
     * 容器后处理-在最后给动态数据源设置可用的数据源
     * @param applicationContext Spring容器上下文
     * @throws BeansException 异常
     */
    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        // 在所有Bean初始化完成后获取动态数据源,最后进行设置,保证自定义的数据源已经被初始化完成了
        DynamicDataSource ds = applicationContext.getBean(DynamicSwitcherConfiguration.DYNAMIC_SOURCE_BEAN_NAME, DynamicDataSource.class);

        // 获取所有的数据源实例(BeanName-DataSource)(需要去除动态数据源-动态数据源也是DataSource类型)
        Map<String, DataSource> beans = applicationContext.getBeansOfType(DataSource.class);
        // 去除自身
        beans.remove(DynamicSwitcherConfiguration.DYNAMIC_SOURCE_BEAN_NAME);

        if (beans.isEmpty()) {
            logger.warn("it seems to you have not setup the dataSource...");
            return;
        }

        // 配置多数据源Map
        Map<Object, Object> dsMap = CollectionUtils.newHashMap(beans.size());
        var hasPrimary = false;
        for (Map.Entry<String, DataSource> entry : beans.entrySet()) {
            var v = entry.getValue();
            var k = entry.getKey();
            // DataSource是在配置类中使用@Bean方法生成的,DataSource类上并没有Primary注解,需要使用Spring容器上下文来获取
            // var p = v.getClass().getAnnotation(Primary.class);
            if (this.isPrimary(applicationContext, k)) {
                if (!hasPrimary) {
                    // 设置主数据源
                    ds.setDefaultTargetDataSource(v);
                    hasPrimary = true;
                    DynamicDataSourceContextHolder.addMap(k, DataSourceType.MASTER, v);
                } else {
                    logger.info("already has a master...{} will become a slave...", k);
                    DynamicDataSourceContextHolder.addMap(k, DataSourceType.SLAVE, v);
                }
            } else {
                DynamicDataSourceContextHolder.addMap(k, DataSourceType.SLAVE, v);
            }
            dsMap.put(k, v);
        }

        // 未指定默认数据源则使用第一个
        if (!hasPrimary) {
            String key = DynamicDataSourceContextHolder.convertSlaveToMaster();
            logger.warn("do not find a primary dataSource...the first dataSource whose name is {} will become the master...", key);
            ds.setDefaultTargetDataSource(beans.get(key));
        }
        // 设置多数据源Map
        // ds.setTargetDataSources(dsMap);
    }

    /**
     * 判断当前数据源是否被Primary注解标注(由于DataSource是在配置类中使用@Bean方法生成的,普通的getClass()再获取注解将不适用)<br/>
     * 而使用容器的context的方法可以获取到
     * @param applicationContext 容器上下文
     * @param key 比对的数据源key
     * @return 是否标注了primary注解
     */
    private boolean isPrimary(ApplicationContext applicationContext, String key) {
        var beanMap = applicationContext.getBeansWithAnnotation(Primary.class);
        return beanMap.containsKey(key);
    }
}
